import operator, time
from typing import Annotated, Any, Sequence
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

def reduce_fanouts(left, right):
    if left is None:
        left = []
    if not right:
        return []
    return left + right

class State(TypedDict):
    aggregate: Annotated[list, operator.add]
    fanout_values: Annotated[list, operator.add]
    which: str

def aggregate_fanout_values(state: State) -> Any:
    ranked_values = sorted(
        state["fanout_values"], key=lambda x: x["reliability"], reverse=True
    )
    return {
        "aggregate":[x["value"] for x in ranked_values] +["I'm E"],
        "fanout_values":[],
    }

def route_bc_or_cd(state: State) -> Sequence[str]:
    if state["which"] == "cd":
        return ["c", "d"]
    return ["b", "c"]

class ReturnNodeValue:
    def __init__(self, node_secret: str):
        self._value = node_secret

    def __call__(self, state: State):
        print(f"Adding {self._value} to {state['aggregate']}")
        return {"aggregate": [self._value]}

class ParallelReturnNodeValue:
    def __init__(self, node_secret: str, reliability: float):
        self._node_secret = node_secret
        self._reliability = reliability

    def __call__(self, state: State):
        print(f"Adding {self._node_secret} to {state['aggregate']} in parallel.")
        return {
            "fanout_values": [
                {
                    "value": [self._node_secret],
                    "reliability": self._reliability
                }
            ]
        }

memory = MemorySaver()

builder = StateGraph(State)
builder.add_node("a", ReturnNodeValue("a"))
builder.add_node("b", ParallelReturnNodeValue("b", 0.9))
builder.add_node("c", ParallelReturnNodeValue("c", 0.1))
builder.add_node("d", ParallelReturnNodeValue("d", 0.3))
builder.add_node("e", aggregate_fanout_values)

intermediates = ["b", "c", "d"]
builder.add_edge(START, "a")
builder.add_conditional_edges("a", route_bc_or_cd, intermediates)
for node in intermediates:
    builder.add_edge(node, "e")
builder.add_edge("e", END)
graph = builder.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "foo"}}
graph.invoke({"aggregate": [], "which": "bc", "fanout_values": []}, config)
print(graph.get_state(config))